# !/usr/bin/env python3
# coding=utf8
"""
sqlite 里的 dbbardata 数据表的保存器(从 BaseDatafeed 把数据查出来, 保存到 dbbardata 数据表)
"""
import argparse
import dataclasses
import datetime
import json
from typing import Dict, List, Set, Tuple, Optional, Union, Callable

from vnpy.trader.constant import Exchange, Interval
from vnpy.trader.database import BaseDatabase, get_database, BarData, BarOverview
from vnpy.trader.datafeed import BaseDatafeed, get_datafeed
from vnpy.trader.utility import extract_vt_symbol

from vnpy_tdxdata.tdxdata_datafeed import TdxdataDatafeed, output

from vnpy_zxtools.chinese_mainland_business_day import ChineseMainlandBusinessDay


class DbBarDataSaver(object):
    """"""

    @dataclasses.dataclass
    class QueryIteem:
        """"""
        symbol: str = ""
        exchange: Exchange = None
        interval: Interval = None
        beg: datetime.datetime = None
        end: datetime.datetime = None

    @classmethod
    def execute(cls, items: List[QueryIteem]) -> None:
        """"""
        database: BaseDatabase = get_database()

        datafeed: BaseDatafeed = get_datafeed()
        assert isinstance(datafeed, TdxdataDatafeed)
        datafeed: TdxdataDatafeed = datafeed

        for item in items:
            data: Optional[List[BarData]] = datafeed.query_bar_history_inner(
                symbol=item.symbol,
                exchange=item.exchange,
                start_business_day=item.beg,
                end_business_day=item.end,
                interval=item.interval,
            )
            data = [] if data is None else data

            if data:
                database.save_bar_data(bars=data)

            message: str = "" \
                f"[{item.exchange.value}.{item.symbol}]\t{item.interval.value}\t" \
                f'[{item.beg.strftime("%Y-%m-%d")} ~ {item.end.strftime("%Y-%m-%d")}] ' \
                f"下载 {len(data)} 条数据"
            output(message)

    @classmethod
    def calculate_query_items_s(
        cls,
        interval_s: List[Interval],
        vt_symbol_s: List[Tuple[str, Exchange]],
        beg_dttm: Optional[datetime.datetime],
        end_dttm: Optional[datetime.datetime],
    ) -> List[QueryIteem]:
        """"""
        if not interval_s:
            # 如果未指定 interval_s 就取默认值,
            interval_s: List[Interval] = [Interval.MINUTE]

        # 因为通达信的数据只有1分钟线、5分钟线，日线, 所以这里仅支持1分钟线和日线,
        supported_interval_s: Tuple[Interval, Interval] = (Interval.MINUTE, Interval.DAILY)

        for interval in interval_s:
            if interval not in supported_interval_s:
                raise ValueError("unsupported Interval '{}'.".format(interval))

        database: BaseDatabase = get_database()

        overviews: List[BarOverview] = database.get_bar_overview()
        overviews = [_ for _ in overviews if _.interval in interval_s]

        if not vt_symbol_s:
            # 如果未指定 vt_symbol_s 就取 BarOverview 数据表里的值,
            vt_symbol_s = sorted(set([(_.symbol, _.exchange) for _ in overviews]))

        default_end: datetime.datetime = datetime.datetime.now()
        default_beg: datetime.datetime = datetime.datetime(year=default_end.year - 1, month=1, day=1)

        items: List[DbBarDataSaver.QueryIteem] = []

        for vt_symbol in vt_symbol_s:
            symbol, exchange = vt_symbol

            for interval in interval_s:
                values = [_ for _ in overviews if (_.symbol, _.exchange, _.interval) == (symbol, exchange, interval)]
                assert len(values) <= 1
                value: BarOverview = values[0] if values else None

                item: DbBarDataSaver.QueryIteem = DbBarDataSaver.QueryIteem(
                    symbol=symbol,
                    exchange=exchange,
                    interval=interval,
                    beg=beg_dttm if beg_dttm else value.end if value else default_beg,
                    end=end_dttm if end_dttm else default_end,
                )
                items.append(item)

        for item in items:
            item.beg = cls.find_bussiness_day(item.beg, is_left=True)
            item.end = cls.find_bussiness_day(item.end, is_left=False)

        for item in items:
            message: str = "" \
                f"[{item.exchange.value}.{item.symbol}]\t{item.interval.value}\t" \
                f'[{item.beg.strftime("%Y-%m-%d")} ~ {item.end.strftime("%Y-%m-%d")}] 即将下载数据'
            output(message)

        return items

    @classmethod
    def find_bussiness_day(cls, base: datetime.datetime, is_left: bool) -> datetime.datetime:
        """"""
        curr: datetime.datetime = datetime.datetime(year=base.year, month=base.month, day=base.day)

        sign: int = -1 if is_left else +1

        is_find: bool = False
        for i in range(90):
            curr += datetime.timedelta(days=(0 if i == 0 else sign * 1))
            if ChineseMainlandBusinessDay.is_business_day(curr):
                is_find: bool = True
                break
        assert is_find

        return curr

    @classmethod
    def calculate_config(cls, args: argparse.Namespace) -> dict:
        """"""
        config: dict = {}
        try:
            if int(args.config) == 1:
                with open(file=__file__ + ".json", mode="r", encoding="utf8") as f:
                    config = json.load(fp=f)
                    assert isinstance(config, dict)
                    if config.get("beg"):
                        config["beg"] = str_to_datetime(config["beg"])
                    if config.get("end"):
                        config["end"] = str_to_datetime(config["end"])
                    if config.get("interval"):
                        config["interval"] = [str_to_interval(_) for _ in config["interval"]]
                    if config.get("vt_symbol"):
                        config["vt_symbol"] = [str_to_tuple_se(_) for _ in config["vt_symbol"]]
        except ValueError:
            config: dict = {}

        if args.beg:
            config["beg"] = args.beg
        if args.end:
            config["end"] = args.end
        if args.vt_symbol:
            config["vt_symbol"] = args.vt_symbol
        if args.interval:
            config["interval"] = args.interval

        return config


def str_to_datetime(src: str) -> datetime.datetime:
    """"""
    try:
        return datetime.datetime.strptime(src, "%Y-%m-%d")
    except ValueError:
        msg = "not a valid 'yyyy-mm-dd' format: '{}'.".format(src)
        raise argparse.ArgumentTypeError(msg)


def str_to_tuple_se(src: str) -> Tuple[str, Exchange]:
    """str => Tuple[symbol, Exchange]"""
    try:
        dst: Tuple[str, Exchange] = extract_vt_symbol(vt_symbol=src)
        return dst
    except ValueError:
        msg = "not a valid vt_symbol(symbol.Exchange) format: '{}'.".format(src)
        raise argparse.ArgumentTypeError(msg)


def str_to_interval(src: str) -> Interval:
    """"""
    try:
        interval: Interval = Interval(src)
        return interval
    except ValueError:
        msg = "not a valid Interval(1m/d) format: '{}'.".format(src)
        raise argparse.ArgumentTypeError(msg)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-b', '--beg', default=None,    type=str_to_datetime)
    parser.add_argument('-e', '--end', default=None,    type=str_to_datetime)
    parser.add_argument('-v', '--vt_symbol', nargs='*', type=str_to_tuple_se)
    parser.add_argument('-i', '--interval', nargs='*',  type=str_to_interval)
    parser.add_argument('-c', '--config', default=None, type=str)
    args = parser.parse_args()

    config: dict = DbBarDataSaver.calculate_config(args=args)

    items: List[DbBarDataSaver.QueryIteem] = DbBarDataSaver.calculate_query_items_s(
        interval_s=config.get("interval"),
        vt_symbol_s=config.get("vt_symbol"),
        beg_dttm=config.get("beg"),
        end_dttm=config.get("end"),
    )

    DbBarDataSaver.execute(items)
